// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/connectors/hive/table_handle.h>
#include <pollux/core/plan_node.h>
#include <pollux/dwio/common/options.h>
#include <pollux/substrait/substrait_to_pollux_expr.h>

namespace kumo::pollux::substrait {
    /// This class is used to convert the Substrait plan into Pollux plan.
    class SubstraitPolluxPlanConverter {
    public:
        explicit SubstraitPolluxPlanConverter(memory::MemoryPool *pool)
            : pool_(pool) {
        }

        struct SplitInfo {
            /// The Partition index.
            u_int32_t partitionIndex;

            /// The file paths to be scanned.
            std::vector<std::string> paths;

            /// The file starts in the scan.
            std::vector<u_int64_t> starts;

            /// The lengths to be scanned.
            std::vector<u_int64_t> lengths;

            /// The file format of the files to be scanned.
            dwio::common::FileFormat format;
        };

        /// Convert Substrait AggregateRel into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::AggregateRel &aggRel);

        /// Convert Substrait ProjectRel into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::ProjectRel &projectRel);

        /// Convert Substrait FilterRel into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::FilterRel &filterRel);

        /// Convert Substrait ReadRel into Pollux PlanNode.
        /// Index: the index of the partition this item belongs to.
        /// Starts: the start positions in byte to read from the items.
        /// Lengths: the lengths in byte to read from the items.
        core::PlanNodePtr toPolluxPlan(
            const ::substrait::ReadRel &readRel,
            std::shared_ptr<SplitInfo> &splitInfo);

        /// Convert Substrait FetchRel into Pollux LimitNode or TopNNode according the
        /// different input of fetchRel.
        core::PlanNodePtr toPolluxPlan(const ::substrait::FetchRel &fetchRel);

        /// Convert Substrait ReadRel into Pollux Values Node.
        core::PlanNodePtr toPolluxPlan(
            const ::substrait::ReadRel &readRel,
            const RowTypePtr &type);

        /// Convert Substrait Rel into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::Rel &rel);

        /// Convert Substrait RelRoot into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::RelRoot &root);

        /// Convert Substrait SortRel into Pollux OrderByNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::SortRel &sortRel);

        /// Convert Substrait Plan into Pollux PlanNode.
        core::PlanNodePtr toPolluxPlan(const ::substrait::Plan &substraitPlan);

        /// Check the Substrait type extension only has one unknown extension.
        bool checkTypeExtension(const ::substrait::Plan &substraitPlan);

        /// Construct the function map between the index and the Substrait function
        /// name.
        void constructFunctionMap(const ::substrait::Plan &substraitPlan);

        /// Return the function map used by this plan converter.
        const melon::F14FastMap<uint64_t, std::string> &getFunctionMap() const {
            return functionMap_;
        }

        /// Return the splitInfo map used by this plan converter.
        const melon::F14FastMap<core::PlanNodeId, std::shared_ptr<SplitInfo> > &
        splitInfos() const {
            return splitInfoMap_;
        }

        /// Looks up a function by ID and returns function name if found. Throws if
        /// function with specified ID doesn't exist. Returns a compound
        /// function specification consisting of the function name and the input
        /// types. The format is as follows: <function
        /// name>:<arg_type0>_<arg_type1>_..._<arg_typeN>
        const std::string &findFunction(uint64_t id) const;

        /// Integrate Substrait emit feature. Here a given 'substrait::RelCommon'
        /// is passed and check if emit is defined for this relation. Basically a
        /// ProjectNode is added on top of 'noEmitNode' to represent output order
        /// specified in 'relCommon::emit'. Return 'noEmitNode' as is
        /// if output order is 'kDriect'.
        core::PlanNodePtr processEmit(
            const ::substrait::RelCommon &relCommon,
            const core::PlanNodePtr &noEmitNode);

    private:
        /// Returns unique ID to use for plan node. Produces sequential numbers
        /// starting from zero.
        std::string nextPlanNodeId();

        /// Used to convert Substrait Filter into Pollux SubfieldFilters which will
        /// be used in TableScan.
        common::SubfieldFilters toPolluxFilter(
            const std::vector<std::string> &inputNameList,
            const std::vector<TypePtr> &inputTypeList,
            const ::substrait::Expression &substraitFilter);

        /// Multiple conditions are connected to a binary tree structure with
        /// the relation key words, including AND, OR, and etc. Currently, only
        /// AND is supported. This function is used to extract all the Substrait
        /// conditions in the binary tree structure into a vector.
        void flattenConditions(
            const ::substrait::Expression &substraitFilter,
            std::vector<::substrait::Expression_ScalarFunction> &scalarFunctions);

        /// The Substrait parser used to convert Substrait representations into
        /// recognizable representations.
        std::shared_ptr<SubstraitParser> substraitParser_{
            std::make_shared<SubstraitParser>()
        };

        /// Helper Function to convert Substrait sortField to Pollux sortingKeys and
        /// sortingOrders.
        std::pair<
            std::vector<core::FieldAccessTypedExprPtr>,
            std::vector<core::SortOrder> >
        processSortField(
            const ::google::protobuf::RepeatedPtrField<::substrait::SortField> &
            sortField,
            const RowTypePtr &inputType);

        /// The Expression converter used to convert Substrait representations into
        /// Pollux expressions.
        std::shared_ptr<SubstraitPolluxExprConverter> exprConverter_;

        /// The unique identification for each PlanNode.
        int planNodeId_ = 0;

        /// The map storing the relations between the function id and the function
        /// name. Will be constructed based on the Substrait representation.
        melon::F14FastMap<uint64_t, std::string> functionMap_;

        /// Mapping from leaf plan node ID to splits.
        melon::F14FastMap<core::PlanNodeId, std::shared_ptr<SplitInfo> >
        splitInfoMap_;

        /// Memory pool.
        memory::MemoryPool *pool_;

        /// Helper function to convert the input of Substrait Rel to Pollux Node.
        template<typename T>
        core::PlanNodePtr convertSingleInput(T rel) {
            POLLUX_CHECK(rel.has_input(), "Child Rel is expected here.");
            return toPolluxPlan(rel.input());
        }
    };
} // namespace kumo::pollux::substrait
